import os
import sys
import requests
from tqdm import tqdm

"""
自动下载SAM2.1 large模型权重的脚本
"""

# SAM 2.1 large模型权重URL
model_url = "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt"
model_filename = "sam2.1_hiera_large.pt"

# 下载目录
download_dir = "../sam2/checkpoints"
save_path = os.path.join(download_dir, model_filename)

def download_file(url, save_path):
    """
    下载文件并显示进度条
    """
    try:
        # 发送GET请求
        response = requests.get(url, stream=True)
        response.raise_for_status()  # 检查请求是否成功
        
        # 获取文件大小
        total_size = int(response.headers.get('content-length', 0))
        
        print(f"开始下载 {model_filename}，文件大小: {total_size/1024/1024:.2f} MB")
        
        # 显示下载进度条
        with open(save_path, 'wb') as file, tqdm(
            desc=os.path.basename(save_path),
            total=total_size,
            unit='iB',
            unit_scale=True,
            unit_divisor=1024,
        ) as bar:
            for data in response.iter_content(chunk_size=8192):
                size = file.write(data)
                bar.update(size)
        
        print(f"成功下载: {save_path}")
        return True
        
    except Exception as e:
        print(f"下载失败: {e}")
        if os.path.exists(save_path):
            os.remove(save_path)  # 清理部分下载的文件
        return False

if __name__ == "__main__":
    print("====================================")
    print("      自动下载SAM2.1 large模型        ")
    print("====================================")
    
    # 创建下载目录
    os.makedirs(download_dir, exist_ok=True)
    
    # 检查文件是否已存在
    if os.path.exists(save_path):
        print(f"文件 {model_filename} 已存在，跳过下载")
    else:
        print(f"开始下载 {model_filename} ({model_url})")
        if not download_file(model_url, save_path):
            print(f"下载 {model_filename} 失败，程序退出")
            sys.exit(1)
    
    print("")
    print("====================================")
    print("       模型权重下载完成！           ")
    print("====================================")